import json
import re
from typing import List, Dict, Any, Optional, Callable, Iterable
import asyncio

import ga_mcp_client
from openai import OpenAI

class FunctionCallWrapper:
    def __init__(self, model_api: Callable, client, tools: Iterable[Any]):
        self.model_api = model_api
        self.client = client
        self.tools = {tool.name: tool for tool in tools}

    def format_parameters(self, schema, indent=2):
        """递归格式化参数结构"""
        if not schema:
            return " " * indent + "(无参数)"
        lines = []
        properties = schema.get("properties", {})
        required = schema.get("required", [])
        for name, prop in properties.items():
            typ = prop.get("type", "object")
            desc = prop.get("description", "")
            req = "（必填）" if name in required else ""
            lines.append(" " * indent + f"- {name}（类型: {typ}{req}）：{desc}")
            # 递归处理嵌套 dto
            if typ == "object" and "properties" in prop:
                lines.append(self.format_parameters(prop, indent + 2))
        return "\n".join(lines) if lines else " " * indent + "(无参数)"

    def get_example_parameters(self, input_schema):
        """生成参数示例，区分 dto 包裹和普通参数"""
        if (
            'properties' in input_schema and
            'dto' in input_schema['properties']
        ):
            # dto 包裹
            dto_props = input_schema['properties']['dto'].get('properties', {})
            example_dto = {k: f"示例值{k}" for k in dto_props.keys()}
            return {"dto": example_dto}
        else:
            props = input_schema.get('properties', {})
            return {k: f"示例值{k}" for k in props.keys()}

    def get_tools_description(self) -> str:
        """生成工具的自然语言描述，包括参数结构和参数示例"""
        description = "可用工具列表：\n"
        for tool in self.tools.values():
            description += f"- {tool.name}: {getattr(tool, 'description', '')}\n"
            input_schema = getattr(tool, 'inputSchema', None) or getattr(tool, 'input_schema', None) or {}
            description += "  参数结构：\n"
            description += self.format_parameters(input_schema, indent=4) + "\n"
            # 参数示例
            example_params = self.get_example_parameters(input_schema)
            description += f"  参数示例：{json.dumps(example_params, ensure_ascii=False)}\n"
        description += "当需要调用工具时，请使用以下格式指定：\n"
        description += "```function_call\n{\n  \"name\": \"工具名称\",\n  \"parameters\": {参数示例}\n}\n```"
        return description

    def generate_prompt(self, user_query: str) -> str:
        tools_description = self.get_tools_description()
        return f"{tools_description}\n\n用户问题：{user_query}\n\n回答："

    def parse_function_call(self, model_response: str) -> Optional[Dict[str, Any]]:
        pattern = r"```function_call\n(.*?)\n```"
        match = re.search(pattern, model_response, re.DOTALL)
        if match:
            try:
                function_call = json.loads(match.group(1))
                return function_call
            except json.JSONDecodeError:
                print("解析工具调用信息失败")
                return None
        return None

    async def run(self, user_query: str) -> Any:
        prompt = self.generate_prompt(user_query)
        
        print(f"生成的提示词：{prompt}")
        
        model_response = self.model_api(prompt)
        # 支持异步模型API（如需）
        if asyncio.iscoroutine(model_response):
            model_response = await model_response
        function_call = self.parse_function_call(model_response)
        if function_call:
            tool_name = function_call.get("name")
            parameters = function_call.get("parameters", {})
            print(f"toolname:{tool_name},parameters{parameters}")
            if tool_name in self.tools:
                try:
                    tool = self.tools[tool_name]
                    input_schema = getattr(tool, 'inputSchema', None) or getattr(tool, 'input_schema', None) or {}
                    # 检查是否需要dto嵌套，且避免重复包裹
                    if (
                        'properties' in input_schema and
                        'dto' in input_schema['properties'] and
                        not (isinstance(parameters, dict) and 'dto' in parameters)
                    ):
                        call_params = {'dto': parameters}
                    else:
                        call_params = parameters
                    result = await self.client.call_tool(tool_name, call_params)
                    
                    print(result)
                    
                    follow_up_prompt = f"工具 {tool_name} 已执行，结果：{getattr(result, 'data', str(result))}\n\n原始问题：{user_query}\n\n最终回答："
                    final_response = self.model_api(follow_up_prompt)
                    if asyncio.iscoroutine(final_response):
                        final_response = await final_response
                    return final_response
                except Exception as e:
                    return f"执行工具 {tool_name} 时出错：{str(e)}"
            else:
                return f"未知工具：{tool_name}"
        else:
            return model_response

def get_model_answer(prompt: str) -> str:
    client = OpenAI(
        api_key="sk-local",
        base_url="http://192.168.210.9:81/v1"
    )
    response = client.chat.completions.create(
        model="TT-14B-R1-0223",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=512,
        temperature=0.2
    )
    return response.choices[0].message.content.strip()

async def main():
    client, tools = await ga_mcp_client.get_tools(ga_mcp_client.MCP_BASE_URL)
    wrapper = FunctionCallWrapper(get_model_answer, client, tools)
    #user_query = "黄文隆的基本信息"
    user_query = "成都市的经纬度"
    # user_query = "查一下黄文隆的轨迹信息"
    response = await wrapper.run(user_query)
    print(response)

if __name__ == "__main__":
    asyncio.run(main())
